from abc import ABC, abstractmethod
import pymysql
import psycopg2


# 数据库连接策略抽象基类
class DatabaseConnectionStrategy(ABC):
    """
    定义了一个所有具体策略类都应该实现的接口。
    """

    @abstractmethod
    def connect(self, data_source):
        """数据库连接"""
        pass

    @abstractmethod
    def get_table_total(self, connection, table_name):
        """获取指定表的数据条数"""
        pass

    @abstractmethod
    def get_table_data(self, connection, table_name, page: int = 1, size : int=500):
        """获取指定表的指定页数据"""
        pass


# MySQL连接策略
class MySQLConnectionStrategy(DatabaseConnectionStrategy):

    def connect(self, data_source):
        """
        建立并返回一个持久的数据库连接。
        :param data_source: 包含连接信息的字典。
        :return: 数据库连接对象。
        """
        try:
            connection = pymysql.connect(
                host=data_source.get('host', 'localhost'),
                user=data_source.get('username'),
                password=data_source.get('password'),
                database=data_source.get('db_name'),
                port=data_source.get('port', 3306),
            )
            return connection
        except Exception as e:
            raise Exception(f"MySQL连接失败: {str(e)}")

    def get_table_total(self, connection, table_name):
        try:
            with connection.cursor() as cursor:
                cursor.execute(f"SELECT COUNT(*) FROM {table_name};")
                result = cursor.fetchone()
                return result[0] if result else 0
        except Exception as e:
            raise Exception(f"获取表 {table_name} 数据总量失败: {str(e)}")

    def get_table_data(self, connection, table_name, page: int = 1, size : int=500):
        try:
            with connection.cursor(pymysql.cursors.DictCursor) as cursor:
                start = (page - 1) * size
                cursor.execute(f"SELECT * FROM {table_name} LIMIT {size} OFFSET {start};")
                result = cursor.fetchall()
                return result
        except Exception as e:
            raise Exception(f"获取表 {table_name} 数据失败: {str(e)}")


# PostgreSQL连接策略
class PostgreSQLConnectionStrategy(DatabaseConnectionStrategy):
    def test_connection(self, data_source):
        try:
            connection = psycopg2.connect(
                host=data_source.get('host', 'localhost'),
                user=data_source.get('username'),
                password=data_source.get('password'),
                dbname=data_source.get('db_name'),
                port=data_source.get('port', 5432),
            )
            with connection.cursor() as cursor:
                cursor.execute("SELECT version();")
                version = cursor.fetchone()
            connection.close()
            return {'message': 'PostgreSQL连接成功', 'version': version}
        except Exception as e:
            raise Exception(f"PostgreSQL连接失败: {str(e)}")


# 定义一个连接器
class DBConnection:
    def __init__(self, strategy: DatabaseConnectionStrategy):
        self._strategy = strategy

    def set_strategy(self, strategy: DatabaseConnectionStrategy):
        self._strategy = strategy

    def test_database_connection(self, data_source):
        return self._strategy.test_connection(data_source)

    def database_connection(self, data_source):
        return self._strategy.connect(data_source)


# 定义一个连接器
class DatabaseService:

    strategy = None
    strategies = {
        'mysql': MySQLConnectionStrategy,
        'postgresql': PostgreSQLConnectionStrategy,
    }
    __connect = None


    def __init__(self, db_type):
        self.strategy = self.strategies.get(db_type)
        self.__connect = None


    # 数据库连接操作，返回数据库连接资源
    def connection(self, data_source):
        if self.strategy:
            self.__connect = self.strategy().connect(data_source)
            return self.__connect
        else:
            raise ValueError("Invalid database type")

    def test_database_connection(self):
        try:
            with self.__connect.cursor() as cursor:
                cursor.execute("SELECT VERSION();")
                version = cursor.fetchone()
            return {'code': 1, 'message': 'MySQL连接成功', 'version': version}
        except Exception as e:
            if self.__connect:
                self.__connect.close()
            raise Exception(f"MySQL连接失败: {str(e)}")
        finally:
            self.__connect.close()

    # 获取指定表的数据条数
    def get_table_rows(self, table_name):
        return self.strategy().get_table_total(self.__connect, table_name)

    # 获取指定表，指定页数据
    def get_table_data(self, table_name, page=1, size=500):
        return self.strategy().get_table_data(self.__connect, table_name, page, size)






